# -*- coding: utf-8 -*-
# !/usr/bin/env python
"""
-------------------------------------------------
   File Name：     datalist
   Description :   
   Author :       lth
   date：          2022/1/24
-------------------------------------------------
   Change Activity:
                   2022/1/24 12:26: create this script
-------------------------------------------------
create this script acted in this project to translate the data from process
"""
__author__ = 'lth'

from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms

from config import GetConfig

train_transform = transforms.Compose(
    [
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5,0.5,0.5], std=[0.5,0.5,0.5])
    ]
)
test_transform = transforms.Compose(
    [
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5,0.5,0.5], std=[0.5,0.5,0.5])
    ]
)

data_transform = transforms.Compose(
    [
        transforms.RandomHorizontalFlip(p=0.5),
        transforms.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.5)
    ]
)


class SRDataset(Dataset):
    def __init__(self, data, mode="train"):
        self.args = GetConfig()

        self.data = data
        self.mode = mode

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        image = Image.open(self.data[index]).convert("RGB")

        width, height = image.size

        if width > height:
            width = self.args.higher_res
            height = int(height / width * self.args.higher_res)
        else:
            height = self.args.higher_res
            width = int(self.args.higher_res / (height / width))

        image = image.resize((width, height), Image.CUBIC)

        new_image = Image.new("RGB", (self.args.higher_res, self.args.higher_res), (128, 128, 128))
        new_image.paste(image, ((self.args.higher_res - width) // 2, (self.args.higher_res - height) // 2))

        new_image = data_transform(new_image)

        image_high = new_image
        image_low = new_image.resize((self.args.lower_res, self.args.lower_res))

        return train_transform(image_low), train_transform(image_high)
